package com.lrj.datastream.sink;

import com.lrj.datastream.domain.ClickEvent;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.api.java.utils.ParameterTool;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.connectors.elasticsearch.ElasticsearchSinkFunction;
import org.apache.flink.streaming.connectors.elasticsearch.RequestIndexer;
import org.apache.flink.streaming.connectors.elasticsearch7.ElasticsearchSink;
import org.apache.http.HttpHost;
import org.elasticsearch.action.index.IndexRequest;
import org.elasticsearch.client.Requests;

import java.util.*;

/**
 * @author lrj
 * @date 2022/3/29 19:20
 */
public class TestElasticSearchSink {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        DataStreamSource<ClickEvent> stream = env.fromElements(
                new ClickEvent("Mary", "./home", 1000L),
                new ClickEvent("Blob", "./home", 2000L),
                new ClickEvent("Mary", "./cart", 3000L)
        );
        ParameterTool parameterTool = ParameterTool.fromArgs(args);

        String host = parameterTool.get("host", "localhost");
        int port = parameterTool.getInt("port", 9200);
        List<HttpHost> hostList = Collections.singletonList(new HttpHost(host, port));
        ElasticsearchSinkFunction<ClickEvent> function = new ElasticsearchSinkFunction<ClickEvent>() {
            private IndexRequest createIndex(String index, Map<String, ?> source) {
                return Requests.indexRequest(index).source(source);
            }

            @Override
            public void process(ClickEvent event, RuntimeContext runtimeContext, RequestIndexer requestIndexer) {
                Map<String, Object> map = new HashMap<>(4);
                map.put("userId", event.getUser());
                map.put("url", event.getUrl());
                map.put("timestamp", event.getTimestamp());
                requestIndexer.add(createIndex("user-"+event.getUser(), map));
            }
        };
        ElasticsearchSink<ClickEvent> sink = new ElasticsearchSink.Builder<>(hostList, function).build();
        stream.addSink(sink);
        env.execute(TestElasticSearchSink.class.getName());
    }


}
